import sys
import torch
from transformers import AutoModelForCausalLM
from transformers import Qwen2_5OmniForConditionalGeneration
import pdb
def compare_models(model_path1, model_path2):
    try:
        model1 = AutoModelForCausalLM.from_pretrained(model_path1, trust_remote_code=True)
        model2 = AutoModelForCausalLM.from_pretrained(model_path2, trust_remote_code=True)
    except:
        model1 = Qwen2_5OmniForConditionalGeneration.from_pretrained(model_path1, trust_remote_code=True)
        model2 = Qwen2_5OmniForConditionalGeneration.from_pretrained(model_path2, trust_remote_code=True)

    # pdb.set_trace()
    model1_params = model1.state_dict()
    model2_params = model2.state_dict()

    del_list = []
    for key in model1_params:
        if 'token2wav' in key or 'talker' in key or 'visual' in key:
            del_list.append(key)
    for key in del_list:
        del model1_params[key]
        del model2_params[key]

    if model1_params.keys() != model2_params.keys():
        print("The models have different architectures.")
        for key in model1_params:
            if key not in model2_params:
                print("model2 missing key:", key)
        for key in model2_params:
            if key not in model1_params:
                print("model1 missing key:", key)
        return

    for key in model1_params:
        sum_diff = torch.sum(model1_params[key]) - torch.sum(model2_params[key])
        diff_sum = torch.sum(torch.abs(model1_params[key] - model2_params[key]))
        diff_avg = diff_sum / model1_params[key].numel()
        if diff_sum.item() > 0.00001 or diff_avg.item() > 0.00001:
            print(f"{key} diff_sum: {diff_sum.item()}")
            # print(f"{key} diff_avg: {diff_avg.item()}")
            # print(f"{key} sum_diff: {sum_diff.item()}")
            # torch.set_printoptions(precision=4, threshold=float('inf'))
            # print(model1_params[key][:, 0])
            # print(model1_params[key][0, :])
            # print("______________________")
            # print(model2_params[key][:, 0])
            # print(model2_params[key][0, :])
            # print("______________________")
            # break
    

if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("Usage: python compare_models.py <model_path1> <model_path2>")
    else:
        compare_models(sys.argv[1], sys.argv[2])
